import torch
from torch import Tensor

    
def mrr_score(pred_batch:Tensor, true_batch:Tensor):
    """Compute mean reciprocal rank (MRR) score.

    Args:
        pred_batch (_type_): _description_(B,C)
        true_batch (_type_): _description_(B)
    """

    rank_pred = torch.argsort(pred_batch, dim=1, descending=True)
    _, hit_pos = torch.where(rank_pred == true_batch.unsqueeze(1))
    
    return torch.mean(1.0 / (hit_pos + 1))

def hit_score(pred_batch:Tensor, true_batch:Tensor, k=5):
    """Compute hit score.

    Args:
        pred_batch (_type_): _description_(B,C)
        true_batch (_type_): _description_(B)
    """

    rank_pred = torch.argsort(pred_batch, dim=1, descending=True)
    hit_mask = (rank_pred == true_batch.unsqueeze(1))[rank_pred < k]
    hit_num = torch.sum(hit_mask)

    return hit_num / pred_batch.size(0)
